import numpy as np
from loguru import logger


class AudioWatermark:
    def __init__(self, watermark_text="MarsX AI", perturbation=0.001):
        """
        初始化水印类。

        :param watermark_text: 要嵌入的水印文本
        :param perturbation: 水印嵌入的扰动大小
        """
        self.watermark_text = watermark_text
        self.watermark_bits = self._text_to_bits(watermark_text)
        self.perturbation = perturbation

    def _text_to_bits(self, text):
        """
        将文本转换为二进制字符串。

        :param text: 输入文本
        :return: 二进制字符串
        """
        return ''.join(format(ord(i), '08b') for i in text)

    def _bits_to_text(self, bits):
        """
        将二进制字符串转换为文本。

        :param bits: 二进制字符串
        :return: 对应的文本
        """
        chars = [chr(int(bits[i:i + 8], 2)) for i in range(0, len(bits), 8)]
        return ''.join(chars)

    def embed_watermark(self, audio, sampling_rate):
        """
        在音频中嵌入水印。如果样本点不足，则不嵌入水印并返回原始音频。

        :param audio: 输入音频的numpy数组
        :param sampling_rate: 采样率
        :return: (采样率, 带水印的音频)
        """
        # 检查音频长度是否足够嵌入水印
        if len(audio) < len(self.watermark_bits):
            logger.warning("音频长度不足，无法嵌入完整水印，返回原始音频。")
            return sampling_rate, audio

        # 创建音频副本，防止修改原始音频
        watermarked_audio = np.copy(audio)

        # 在音频中插入水印，简单通过正负扰动
        for i, bit in enumerate(self.watermark_bits):
            if bit == '1':
                watermarked_audio[i] += self.perturbation  # 轻微正扰动
            else:
                watermarked_audio[i] -= self.perturbation  # 轻微负扰动

        return sampling_rate, watermarked_audio

    def detect_watermark(self, audio, watermark_length):
        """
        从音频中检测水印。

        :param audio: 输入音频的numpy数组
        :param watermark_length: 要检测的水印长度（字符数）
        :return: 检测到的水印文本
        """
        detected_bits = []

        # 假设水印嵌入在音频的前面几个样本点
        for i in range(watermark_length * 8):  # 每个字符对应8位二进制数据
            if i >= len(audio):
                logger.warning("音频长度不足，无法检测到完整水印。")
                break

            if audio[i] > 0:
                detected_bits.append('1')
            else:
                detected_bits.append('0')

        # 将二进制位还原为水印文本
        detected_watermark = self._bits_to_text(''.join(detected_bits))

        return detected_watermark

# 测试本地音频文件检测水印
if __name__ == '__main__':
    import librosa
    # 创建音频水印对象
    audio_watermark = AudioWatermark()

    # 读取音频文件
    audio_file = r"C:\Users\Adminis\Desktop\20240919_0919007_demo.wav"
    audio, sampling_rate = librosa.load(audio_file, sr=None)

    # 检测水印
    watermark_length = len(audio_watermark.watermark_text)
    detected_watermark = audio_watermark.detect_watermark(audio, watermark_length)
    logger.info(f"检测到的水印文本: {detected_watermark}")

